# @Author : wangyuchen
# @Time : 2021-05-13 16:29

from torch.utils import data
import os
import pandas as pd
from PIL import Image


def pil_loader(path):
    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')


def accimage_loader(path):
    import accimage
    try:
        return accimage.Image(path)
    except IOError:
        # Potentially a decoding problem, fall back to PIL.Image
        return pil_loader(path)


def default_loader(path):
    from torchvision import get_image_backend
    if get_image_backend() == 'accimage':
        return accimage_loader(path)
    else:
        return pil_loader(path)


class TrainDataset(data.Dataset):
    def __init__(self, img_dir, transform=None, loader=default_loader):
        self.img_labels = {'name': [], 'label': []}
        self.img_dir = img_dir
        self.transform = transform
        self.loader = loader
        for rt, dirs, files in os.walk(self.img_dir):
            for f in files:
                if f not in ['.DS_Store']:
                    full_name = os.path.join(rt, f)
                    self.img_labels['name'].append(full_name)
                    self.img_labels['label'].append(rt.split('/')[-1])
        self.img_labels = pd.DataFrame(data=self.img_labels)

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = self.img_labels.iloc[idx]['name']
        image = self.loader(img_path)
        label = self.img_labels.iloc[idx]['label']
        # print(label)
        if self.transform:
            image = self.transform(image)
        return image, label


class ValDataset(data.Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, loader=default_loader):
        self.img_labels = pd.read_csv(annotations_file, sep=' ')
        self.img_dir = img_dir
        self.transform = transform
        self.loader = loader

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = self.loader(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        return self.img_labels.iloc[idx, 0], image, label
        # return image, label


class TestDataset(data.Dataset):
    def __init__(self, img_dir, transform=None, target_transform=None, loader=default_loader):
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader
        self.filenames = []
        for rt, dirs, files in os.walk(self.img_dir):
            for f in files:
                self.filenames.append(f)
        print(self.filenames)

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.filenames[idx])
        image = self.loader(img_path)
        if self.transform:
            image = self.transform(image)
        return image, self.filenames[idx]
